# https://www.runoob.com/w3cnote/python-lzw.html

class LZW:
    
    def compress(self, uncompressed):
        # 创建字典，初始化为256个ASCII字符
        dictionary = {chr(i):i for i in range(0,256)}
 
        p = ''
        result = []
        
        for c in uncompressed:
            pc = p + c
            if pc in dictionary:
                p = pc
            else:
                result.append(dictionary[p])
                dictionary[pc] = len(dictionary)
                p = c
        
        if p != '':
            result.append(dictionary[p])
        return result

    def decompress(self, compressed):
        # 创建字典
        dictionary = {i:chr(i) for i in range(0,256)}
        
        result = []
        p = compressed.pop(0)
        result.append(dictionary[p])
        
        for c in compressed:
            if c in dictionary:
                entry = dictionary[c]
            result.append(entry)
            dictionary[len(dictionary)] = dictionary[p] + entry[0]
            p = c
        return "".join(result)
    
class LZWByte:
    
    def compress(self, uncompressed):
        # 创建字典，初始化为256个ASCII字符
        dictionary = {chr(i):i for i in range(0,256)}
 
        p = ''
        result = []
        
        for c in uncompressed:
            pc = p + chr(c)
            if pc in dictionary:
                p = pc
            else:
                result.append(dictionary[p])
                dictionary[pc] = len(dictionary)
                p = chr(c)
        
        if p != '':
            result.append(dictionary[p])
        return result

    def decompress(self, compressed):
        # 创建字典
        dictionary = {i:chr(i) for i in range(0,256)}
        
        result = []
        p = compressed.pop(0)
        result.append(dictionary[p])
        
        for c in compressed:
            if c in dictionary:
                entry = dictionary[c]
            result.append(entry)
            dictionary[len(dictionary)] = dictionary[p] + entry[0]
            p = c
        return "".join(result).encode(encoding='utf-8')
    

class LZWGif:
    
    def compress(self, uncompressed):
        # 创建字典，初始化为256个ASCII字符
        dictionary = {chr(i):i for i in range(0,256)}
 
        p = ''
        result = []
        
        for c in uncompressed:
            pc = p + chr(c)
            if pc in dictionary:
                p = pc
            else:
                result.append(dictionary[p])
                dictionary[pc] = len(dictionary)
                p = chr(c)
        
        if p != '':
            result.append(dictionary[p])
        return result

    def decompress(self, compressed):
        data = list(compressed)
        # 创建字典
        dictionary = {i:chr(i) for i in range(0,256)}
        
        result = []
        p = data.pop(0)
        result.append(dictionary[p])
        
        for c in data:
            if c in dictionary:
                entry = dictionary[c]
            result.append(entry)
            dictionary[len(dictionary)] = dictionary[p] + entry[0]
            p = c
        return "".join(result).encode(encoding='utf-8')

# 使用示例
# lzw = LZW()
# inputStr = "1234123412341234"
# print("Original:", inputStr, "length:", len(inputStr))
# compressed_data = lzw.compress(inputStr)
# print("Compressed:", compressed_data, "length:", len(compressed_data))
# decompressed_data = lzw.decompress(compressed_data)
# print("Decompressed:", decompressed_data, "length:", len(decompressed_data))

# # lzw = LZWByte()
# inputStr = bytes([65,66,67,65,66,67,65,66,67])
# # print("Original:", inputStr, "length:", len(inputStr))
# # compressed_data = lzw.compress(inputStr)
# # print("Compressed:", compressed_data, "length:", len(compressed_data))
# # decompressed_data = lzw.decompress(compressed_data)
# # print("Decompressed:", decompressed_data, "length:", len(decompressed_data))
